#!/usr/bin/env python

import os, sys
import numpy as np
from astropy.io import fits
from scipy.signal import savgol_filter
from scipy.stats import trim_mean, sigmaclip
import matplotlib.pylab as plt

# CCD data configuration
class ccdDataConfig:
    def __init__(self):
        self.raw_nx = 1250      # number of cols of each OS in the raw fits
        self.raw_ny = 4700      # number of rows of each OS in the raw fits
        self.pre_nx = 27        # number of horizontal prescans of each OS in the raw fits
        self.pre_ny = 0         # number of vertical prescan of each OS in the raw fits
        self.nx     = 1152      # number of real pixels in a row of each OS
        self.ny     = 4616      # number of real pixel rows of each OS
        self.over_nx= self.raw_nx - (self.pre_nx + self.nx)
        self.over_ny= self.raw_ny - (self.pre_ny + self.ny)
        
    def check_img_size(self,img):
        if img.shape[0] != self.raw_ny*2 or img.shape[1] != self.raw_nx*8:
            print('img.shape should be: [{}]'.format(self.raw_ny*2,self.raw_nx*8))
            sys.exit(0)
        
    def get_col_overscan(self, img, ch=0):
        self.check_img_size(img)
        if ch < 4:
        # OS:1~4
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx + self.pre_nx + self.nx
            idx_col2 = idx_col1 + self.over_nx
        elif ch < 8:
        # OS:5~8
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx
            idx_col2 = idx_col1 + self.over_nx
        elif ch < 12:
        # OS:9~12
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx
            idx_col2 = idx_col1 + self.over_nx
        else:
        # OS:13~16
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx + self.pre_nx + self.nx
            idx_col2 = idx_col1 + self.over_nx
            
#        print('col_overscan [ch {}]: '.format(ch),idx_row1, idx_row2, idx_col1, idx_col2)
        return img[idx_row1:idx_row2, idx_col1:idx_col2]
        
    def get_row_overscan(self, img, ch=0):
        self.check_img_size(img)
        if ch < 4:
        # OS:1~4
            idx_row1 = self.ny
            idx_row2 = idx_row1 + self.over_ny
            idx_col1 = ch*self.raw_nx + self.pre_nx
            idx_col2 = idx_col1 + self.nx
        elif ch < 8:
        # OS:5~8
            idx_row1 = self.ny
            idx_row2 = idx_row1 + self.over_ny
            idx_col1 = ch*self.raw_nx + self.over_nx
            idx_col2 = idx_col1 + self.nx
        elif ch < 12:
        # OS:9~12
            idx_row1 = self.raw_ny
            idx_row2 = idx_row1 + self.over_ny
            idx_col1 = (15-ch)*self.raw_nx + self.over_nx
            idx_col2 = idx_col1 + self.nx
        else:
        # OS:13~16
            idx_row1 = self.raw_ny
            idx_row2 = idx_row1 + self.over_ny
            idx_col1 = (15-ch)*self.raw_nx + self.pre_nx
            idx_col2 = idx_col1 + self.nx
        
#        print('row_overscan: ', idx_row1, idx_row2, idx_col1, idx_col2)
        return img[idx_row1:idx_row2, idx_col1:idx_col2]
        
    def get_col_prescan(self, img, ch=0):
        self.check_img_size(img)
        if ch < 4:
        # OS:1~4
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx
            idx_col2 = idx_col1 + self.pre_nx
        elif ch < 8:
        # OS:5~8
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx + self.over_nx + self.nx
            idx_col2 = idx_col1 + self.pre_nx
        elif ch < 12:
        # OS:9~12
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx + self.over_nx + self.nx
            idx_col2 = idx_col1 + self.pre_nx
        else:
        # OS:13~16
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx
            idx_col2 = idx_col1 + self.pre_nx
            
#        print('col_prescan: ', idx_row1, idx_row2, idx_col1, idx_col2)
        return img[idx_row1:idx_row2, idx_col1:idx_col2]

    def get_img_chan(self, img, ch=0):
        self.check_img_size(img)
        if ch < 4:
        # OS:1~4
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx + self.pre_nx
            idx_col2 = idx_col1 + self.nx
        elif ch < 8:
        # OS:5~8
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.raw_nx + self.over_nx
            idx_col2 = idx_col1 + self.nx
        elif ch < 12:
        # OS:9~12
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx + self.over_nx
            idx_col2 = idx_col1 + self.nx
        else:
        # OS:13~16
            idx_row1 = self.raw_ny + self.over_ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.raw_nx + self.pre_nx
            idx_col2 = idx_col1 + self.nx
        
#        print('row_overscan: ', idx_row1, idx_row2, idx_col1, idx_col2)
        return img[idx_row1:idx_row2, idx_col1:idx_col2]
        
    def getImgChanRegion(self, ch=0):
        if ch < 8:
            idx_row1 = 0
            idx_row2 = idx_row1 + self.ny
            idx_col1 = ch*self.nx
            idx_col2 = idx_col1 + self.nx
        else:
            idx_row1 = self.ny
            idx_row2 = idx_row1 + self.ny
            idx_col1 = (15-ch)*self.nx
            idx_col2 = idx_col1 + self.nx
        
        return idx_row1, idx_row2, idx_col1, idx_col2

    def removePreOverScan(self,img):   # for debug
        d = np.zeros((self.ny*2, self.nx*8), dtype=np.float32)
        print(d.shape)
 
        for i in range(16):
            ch_data = self.get_img_chan(img=img, ch=i)
#            print('ch_data.shape: {}'.format(ch_data.shape))
            if i < 8:
                idx_row1 = 0
                idx_row2 = idx_row1 + self.ny
                idx_col1 = i*self.nx
                idx_col2 = idx_col1 + self.nx
            else:
                idx_row1 = self.ny
                idx_row2 = idx_row1 + self.ny
                idx_col1 = (15-i)*self.nx
                idx_col2 = idx_col1 + self.nx

#            print(idx_row1, idx_row2, idx_col1, idx_col2)
            d[idx_row1:idx_row2, idx_col1:idx_col2] = ch_data.astype(np.float32)
#        
        hdu = fits.PrimaryHDU(d)
        hdu.writeto('xx.fits', overwrite=True)


    def runTest(self):
#        for i in range(16):
#            self.get_col_overscan(ch=i)
            
#        for ch in range(16):
#            self.get_row_overscan(ch=ch)
            
#        for ch in range(16):
#            self.get_col_prescan(ch=ch)

#        for ch in range(16):
#            self.get_img_chan(ch=ch)

        img = fits.getdata(sys.argv[1]).astype(np.uint16)
        self.removePreOverScan(img=img)

# pre/overscan data filter parameter settings
class ccdDataFilter:
    def __init__(self):
        self.filterWidth_v  = 250
        self.filterWidth_h  = 5
        
    def update(self, fW_v=250, fW_h=5):
        self.filterWidth_v = fW_v
        self.filterWidth_h = fW_h

    def smooth_data(self, data, axis=0):
        dd = savgol_filter(data, filterwidth * 4 + 1, 3)
        diff = data - dd
        index = np.abs(diff) > sigmaclip(diff, 5, 5)[0].std() * 5
        x = np.arange(len(data))
        dd = np.interp(x, x[~index], data[~index])
        smoothdata = savgol_filter(dd, filterwidth * 2 + 1, 2)
        return smoothdata

def _smooth_data(data, filterwidth):
    dd = savgol_filter(data, filterwidth * 4 + 1, 3)
    diff = data - dd
    index = np.abs(diff) > sigmaclip(diff, 5, 5)[0].std() * 5
    x = np.arange(len(data))
    dd = np.interp(x, x[~index], data[~index])
    smoothdata = savgol_filter(dd, filterwidth * 2 + 1, 2)
    return smoothdata

def ccdPreOverScanCorrection(flatFname, applyGainCorrection=False):
    """
    flatFname: fits filename of a flat
    applyGainCorrection: if True, then apply gain correction using row overscan
    """
    cdc = ccdDataConfig()
    cdf = ccdDataFilter()

    ro_gain = None
    if applyGainCorrection:
        ro_gain = []

    img = fits.getdata(flatFname).astype(np.float32)
    img_corr = np.zeros((9232,9216), dtype=np.float32)
    for ch in range(16):
        ch_data = cdc.get_img_chan(img=img, ch=ch)

        img_corr_tmp = ch_data.astype(np.float32)
        
        # 1) correct row overscan
        print("=> correcting row overscan ...")
        if ch < 8:
            row_over= cdc.get_row_overscan(img=img, ch=ch)[5:,:]    # remove first 5 rows
        else:
            row_over= cdc.get_row_overscan(img=img, ch=ch)[:-5,:]    # remove first 5 rows

        row_over_trim_mean = trim_mean(row_over, 0.1, axis=0)
        row_over_trim_mean_smoothed = _smooth_data(row_over_trim_mean, filterwidth=cdf.filterWidth_h)
        
        for i in range(img_corr_tmp.shape[0]):
            img_corr_tmp[i,:] -= row_over_trim_mean_smoothed

#        # 2) correct col prescan
        print("+> correcting col prescan ...")
        col_pre = cdc.get_col_prescan(img=img, ch=ch)
        col_pre_trim_mean = trim_mean(col_pre, 0.1, axis=1)
        col_pre_trim_mean_smoothed = _smooth_data(col_pre_trim_mean, filterwidth=cdf.filterWidth_v)

        for j in range(img_corr_tmp.shape[1]):
            img_corr_tmp[:,j] -= (col_pre_trim_mean_smoothed-np.median(col_pre_trim_mean_smoothed))

        if applyGainCorrection:
            gain = row_over_trim_mean_smoothed / row_over_trim_mean_smoothed[0]
            ro_gain.append(gain)

            for i in range(img_corr_tmp.shape[0]):
                img_corr_tmp[i,:] = img_corr_tmp[i,:] / gain

        idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)
        img_corr[idx_row1:idx_row2, idx_col1:idx_col2] = img_corr_tmp

#    hdu = fits.PrimaryHDU(img_corr)
#    hdu.writeto('ffkk.fits', overwrite=True)
#    cmd = 'ds9 {}&'.format('ffkk.fits')
#    os.system(cmd)
    if applyGainCorrection:
        ro_gain = np.array(ro_gain)

    return img_corr, ro_gain

def ccdOverScanCorrection(flatFname, applyGainCorrection=False):
    """
    modifiled from ccdPreOverScanCorrection(flatFname, applyGainCorrection=False), use col-overscan instead of col-prescan
    flatFname: fits filename of a flat
    applyGainCorrection: if True, then apply gain correction using row overscan
    """
    cdc = ccdDataConfig()
    cdf = ccdDataFilter()

    ro_gain = None
    if applyGainCorrection:
        ro_gain = []

    img = fits.getdata(flatFname).astype(np.float32)
    img_corr = np.zeros((9232,9216), dtype=np.float32)
    for ch in range(16):
        ch_data = cdc.get_img_chan(img=img, ch=ch)

        img_corr_tmp = ch_data.astype(np.float32)

        # 1) correct row overscan
        print("=> correcting row overscan ...")
        if ch < 8:
            row_over= cdc.get_row_overscan(img=img, ch=ch)[5:,:]    # remove first 5 rows
        else:
            row_over= cdc.get_row_overscan(img=img, ch=ch)[:-5,:]    # remove first 5 rows

        row_over_trim_mean = trim_mean(row_over, 0.1, axis=0)
        row_over_trim_mean_smoothed = _smooth_data(row_over_trim_mean, filterwidth=cdf.filterWidth_h)

        for i in range(img_corr_tmp.shape[0]):
            img_corr_tmp[i,:] -= row_over_trim_mean_smoothed

#        # 2) correct col overscan
        print("+> correcting col overscan ...")
        col_over = cdc.get_col_overscan(img=img, ch=ch)
        col_over_trim_mean = trim_mean(col_over, 0.1, axis=1)
        col_over_trim_mean_smoothed = _smooth_data(col_over_trim_mean, filterwidth=cdf.filterWidth_v)

        for j in range(img_corr_tmp.shape[1]):
            img_corr_tmp[:,j] -= (col_over_trim_mean_smoothed-np.median(col_over_trim_mean_smoothed))

        if applyGainCorrection:
            gain = row_over_trim_mean_smoothed / row_over_trim_mean_smoothed[0]
            ro_gain.append(gain)

            for i in range(img_corr_tmp.shape[0]):
                img_corr_tmp[i,:] = img_corr_tmp[i,:] / gain

        idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)
        img_corr[idx_row1:idx_row2, idx_col1:idx_col2] = img_corr_tmp

#    hdu = fits.PrimaryHDU(img_corr)
#    hdu.writeto('ffkk.fits', overwrite=True)
#    cmd = 'ds9 {}&'.format('ffkk.fits')
#    os.system(cmd)
    if applyGainCorrection:
        ro_gain = np.array(ro_gain)

    return img_corr, ro_gain

def estimateGainUsingCorrectedFlats(corr_flat1, corr_flat2, sClip=False, applyFlatCorrection=False):
    """
    sigmaClip range parameters are changed to (3,3), using nSigma=3
    """
    nSigma = 3

    cdc = ccdDataConfig()
    cdf = ccdDataFilter()
    
    cf1 = corr_flat1.copy()
    cf2 = corr_flat2.copy()

    flat_mean = []
    flat_var = []
    gain = []


    if applyFlatCorrection:
        for ch in range(16):
            idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)
            flatProfile1 = np.median( cf1[idx_row1:idx_row2, idx_col1:idx_col2], axis=0 )
            flatProfile2 = np.median( cf2[idx_row1:idx_row2, idx_col1:idx_col2], axis=0 )
            for i in range( idx_row1, idx_row2 ):
                cf1[i,idx_col1:idx_col2] = cf1[i,idx_col1:idx_col2] / (flatProfile1) * np.median(flatProfile1)
                cf2[i,idx_col1:idx_col2] = cf2[i,idx_col1:idx_col2] / (flatProfile2) * np.median(flatProfile2)

    fmean = ( cf1.astype(np.float32) + cf2.astype(np.float32) ) / 2
    fdiff = ( cf1.astype(np.float32) - cf2.astype(np.float32) )

    for ch in range(16):
        idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)

        if sClip:
            tmp_mean = np.median( sigmaclip(fmean[idx_row1:idx_row2, idx_col1:idx_col2].flatten(), nSigma, nSigma)[0] )
            tmp_var  = np.var( sigmaclip(fdiff[idx_row1:idx_row2, idx_col1:idx_col2].flatten(), nSigma, nSigma)[0] ) / 2
        else:
            tmp_mean = np.median( fmean[idx_row1:idx_row2, idx_col1:idx_col2] )
            tmp_var  = np.var( fdiff[idx_row1:idx_row2, idx_col1:idx_col2] ) / 2

        flat_mean.append( tmp_mean )
        flat_var.append( tmp_var )
        gain.append( tmp_mean / tmp_var )

    return np.array(flat_mean), np.array(flat_var), np.array(gain)


if __name__ == '__main__':
#    cdc = ccdDataConfig()
#    cdc.runTest()

    ccdPreOverScanCorrection(flatFname=sys.argv[1], biasFname=None)
    
